import os
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset


# 自定义数据集类
class KneeDataset(Dataset):
    def __init__(self, image_dir, label_dir, transform=None):
        self.image_paths = sorted([os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith('.jpg')])
        self.label_paths = sorted([os.path.join(label_dir, f) for f in os.listdir(label_dir) if f.endswith('.png')])
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        # 加载图像和标注
        img = cv2.imread(self.image_paths[idx], cv2.IMREAD_GRAYSCALE)
        label = cv2.imread(self.label_paths[idx], cv2.IMREAD_GRAYSCALE)

        # 数据预处理
        img = cv2.resize(img, (512, 512)) / 255.0  # 归一化
        label = cv2.resize(label, (512, 512), interpolation=cv2.INTER_NEAREST)
        label = (label > 0).astype(np.float32)

        # 转换为PyTorch张量
        img = torch.tensor(img, dtype=torch.float32).unsqueeze(0)  # 添加通道维度
        label = torch.tensor(label, dtype=torch.float32).unsqueeze(0)

        return img, label
